import json
import os

import click
import dotenv
import pandas as pd
import requests
from tqdm import tqdm


@click.command()
@click.argument("file", type=click.Path(exists=True))
@click.option("--dataset-id", help="dataset id", required=True)
@click.option("--document-id", help="document id", required=True)
@click.option("--sheet-name", help="sheet name", default=None)
def main(file, dataset_id, document_id, sheet_name):
    """Upload dataset to Dify from Excel file."""
    dotenv.load_dotenv()
    api_key = os.getenv("DIFY_API_KEY_DB")
    base_url = os.getenv("DIFY_BASE_URL")

    if not sheet_name:
        df = pd.read_excel(file)
    else:
        df = pd.read_excel(file, sheet_name=sheet_name)

    # fillna
    df = df.fillna("")
    df = df.astype(str)

    for _, row in tqdm(df.iterrows()):
        url = f"{base_url}/v1/datasets/{dataset_id}/documents/{document_id}/segments"
        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        }

        data = {
            "segments": [
                {
                    "content": f"{row['一级分类']}\n{row['二级分类']}\n{row['岗位名称']}",
                    "answer": json.dumps(
                        row.to_dict(), ensure_ascii=False, indent=4
                    ),
                    "keywords": [row["一级分类"], row["二级分类"], row["岗位名称"]],
                }
            ]
        }
        response = requests.post(url, headers=headers, json=data)
        response.raise_for_status()


if __name__ == "__main__":
    main()
